#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>
#include <sys/time.h>

const int TILE_WIDTH = 16;

static inline double get_cpu_time_ms() {
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return (ts.tv_sec * 1000.0) + (ts.tv_nsec / 1000000.0);
}

static void ConvCPU(const double* img, const double* ker, double* out, int Row, int Col, int K) {
	int mf = Row - K + 1;
	int nf = Col - K + 1;
	for (int i = 0; i < mf; i++) {
		for (int j = 0; j < nf; j++) {
			double sum = 0.0;
			for (int p = 0; p < K; p++) {
				for (int q = 0; q < K; q++) {
					sum += img[(i + p) * Col + (j + q)] * ker[p * K + q];
				}
			}
			out[i * nf + j] = sum;
		}
	}
}

__global__ void ConvKernel(const double* __restrict__ img,
						   const double* __restrict__ ker,
						   double* __restrict__ out,
						   int Row, int Col, int K) {
	int j = blockIdx.x * blockDim.x + threadIdx.x;
	int i = blockIdx.y * blockDim.y + threadIdx.y;
	int mf = Row - K + 1;
	int nf = Col - K + 1;
	if (i >= mf || j >= nf) return;
	double sum = 0.0;
	for (int p = 0; p < K; p++) {
		for (int q = 0; q < K; q++) {
			sum += img[(i + p) * Col + (j + q)] * ker[p * K + q];
		}
	}
	out[i * nf + j] = sum;
}

__global__ void ConvKernelShared(const double* __restrict__ img,
								 const double* __restrict__ ker,
								 double* __restrict__ out,
								 int Row, int Col, int K) {
	extern __shared__ double sh[];
	int j0 = blockIdx.x * blockDim.x;
	int i0 = blockIdx.y * blockDim.y;
	int mf = Row - K + 1;
	int nf = Col - K + 1;
	if (i0 >= mf || j0 >= nf) return;

	int tile_out_h = min(blockDim.y, mf - i0);
	int tile_out_w = min(blockDim.x, nf - j0);
	int H = tile_out_h + K - 1;
	int W = tile_out_w + K - 1;

	for (int y = threadIdx.y; y < H; y += blockDim.y) {
		int gy = i0 + y;
		for (int x = threadIdx.x; x < W; x += blockDim.x) {
			int gx = j0 + x;
			sh[y * W + x] = img[gy * Col + gx];
		}
	}
	__syncthreads();

	int tx = threadIdx.x;
	int ty = threadIdx.y;
	if (ty < tile_out_h && tx < tile_out_w) {
		double sum = 0.0;
		for (int p = 0; p < K; ++p) {
			int sy = ty + p;
			int srow = sy * W;
			for (int q = 0; q < K; ++q) {
				sum += sh[srow + (tx + q)] * ker[p * K + q];
			}
		}
		int i = i0 + ty;
		int j = j0 + tx;
		out[i * nf + j] = sum;
	}
}

static void verify(const double* ref, const double* out, size_t len) {
	size_t err_cnt = 0;
	const double eps = 1e-9;
	for (size_t i = 0; i < len; ++i) {
		double e = fabs(ref[i] - out[i]);
		if (e > eps) err_cnt++;
	}
	if (err_cnt) {
		printf("Verification failed! Number of mismatches: %zu\n", err_cnt);
	} else {
		printf("Verification passed!\n");
	}
}

int main() {
	const char* inputPath = "input3.txt";
	const char* outputPath = "output3.txt";

	int Row = 0, Col = 0, K = 0;
	FILE* fin = fopen(inputPath, "r");
	fscanf(fin, "%d,%d,%d", &Row, &Col, &K);
	fclose(fin);

	size_t imgSize = (size_t)Row * (size_t)Col;
	size_t kerSize = (size_t)K * (size_t)K;
	int mf = Row - K + 1;
	int nf = Col - K + 1;
	if (mf <= 0 || nf <= 0) {
		fprintf(stderr, "Invalid sizes for valid convolution\n");
		return 1;
	}

	double* h_img = (double*)malloc(imgSize * sizeof(double));
	double* h_ker = (double*)malloc(kerSize * sizeof(double));
	double* h_out_cpu = (double*)malloc((size_t)mf * (size_t)nf * sizeof(double));
	double* h_out_gpu = (double*)malloc((size_t)mf * (size_t)nf * sizeof(double));
	if (!h_img || !h_ker || !h_out_cpu || !h_out_gpu) {
		fprintf(stderr, "Host memory allocation failed\n");
		return 1;
	}

	srand(0xC0FFEE);
	for (size_t i = 0; i < imgSize; ++i) {
		h_img[i] = (double)(rand() % 10);
	}
	for (size_t i = 0; i < kerSize; ++i) {
		h_ker[i] = (double)((rand() % 5) - 2);
	}

	double cpu_start = get_cpu_time_ms();
	ConvCPU(h_img, h_ker, h_out_cpu, Row, Col, K);
	double cpu_stop = get_cpu_time_ms();
	double cpu_time = cpu_stop - cpu_start;
	printf("CPU Time: %.2f ms\n", cpu_time);

	double *d_img = nullptr, *d_ker = nullptr, *d_out = nullptr;
	cudaError_t err;
	err = cudaMalloc((void**)&d_img, imgSize * sizeof(double)); if (err) { fprintf(stderr, "cudaMalloc d_img failed: %s\n", cudaGetErrorString(err)); return 1; }
	err = cudaMalloc((void**)&d_ker, kerSize * sizeof(double)); if (err) { fprintf(stderr, "cudaMalloc d_ker failed: %s\n", cudaGetErrorString(err)); return 1; }
	err = cudaMalloc((void**)&d_out, (size_t)mf * (size_t)nf * sizeof(double)); if (err) { fprintf(stderr, "cudaMalloc d_out failed: %s\n", cudaGetErrorString(err)); return 1; }

	cudaMemcpy(d_img, h_img, imgSize * sizeof(double), cudaMemcpyHostToDevice);
	cudaMemcpy(d_ker, h_ker, kerSize * sizeof(double), cudaMemcpyHostToDevice);

	dim3 block(TILE_WIDTH, TILE_WIDTH);
	dim3 grid((nf + TILE_WIDTH - 1) / TILE_WIDTH, (mf + TILE_WIDTH - 1) / TILE_WIDTH);

	cudaEvent_t startEvent, stopEvent;
	cudaEventCreate(&startEvent);
	cudaEventCreate(&stopEvent);

	cudaMemset(d_out, 0, (size_t)mf * (size_t)nf * sizeof(double));
	cudaEventRecord(startEvent);
	ConvKernel<<<grid, block>>>(d_img, d_ker, d_out, Row, Col, K);
	cudaEventRecord(stopEvent);
	cudaEventSynchronize(stopEvent);

	float elapsedMs = 0.0f;
	cudaEventElapsedTime(&elapsedMs, startEvent, stopEvent);
	double gpu_time_naive = (double)elapsedMs;
	printf("GPU Time (naive): %.2f ms\n", gpu_time_naive);

	cudaMemcpy(h_out_gpu, d_out, (size_t)mf * (size_t)nf * sizeof(double), cudaMemcpyDeviceToHost);

	verify(h_out_cpu, h_out_gpu, (size_t)mf * (size_t)nf);

	double *d_out_shared = nullptr;
	cudaMalloc((void**)&d_out_shared, (size_t)mf * (size_t)nf * sizeof(double));
	cudaMemset(d_out_shared, 0, (size_t)mf * (size_t)nf * sizeof(double));
	double* h_out_gpu_shared = (double*)malloc((size_t)mf * (size_t)nf * sizeof(double));
	size_t smem_bytes = (size_t)(TILE_WIDTH + K - 1) * (size_t)(TILE_WIDTH + K - 1) * sizeof(double);

	cudaEventRecord(startEvent);
	ConvKernelShared<<<grid, block, smem_bytes>>>(d_img, d_ker, d_out_shared, Row, Col, K);
	cudaEventRecord(stopEvent);
	cudaEventSynchronize(stopEvent);
	float elapsedMsShared = 0.0f;
	cudaEventElapsedTime(&elapsedMsShared, startEvent, stopEvent);
	double gpu_time_shared = (double)elapsedMsShared;
	printf("GPU Time (shared): %.2f ms\n", gpu_time_shared);
	cudaMemcpy(h_out_gpu_shared, d_out_shared, (size_t)mf * (size_t)nf * sizeof(double), cudaMemcpyDeviceToHost);
	verify(h_out_cpu, h_out_gpu_shared, (size_t)mf * (size_t)nf);

	FILE* fout = fopen(outputPath, "w");
	if (fout) {
		fprintf(fout, "%.2f,%.2f,%.2f", cpu_time, gpu_time_naive, gpu_time_shared);
		fclose(fout);
	} else {
		fprintf(stderr, "Failed to open %s for writing\n", outputPath);
	}

	cudaEventDestroy(startEvent);
	cudaEventDestroy(stopEvent);
	cudaFree(d_img);
	cudaFree(d_ker);
	cudaFree(d_out);
	cudaFree(d_out_shared);
	free(h_img);
	free(h_ker);
	free(h_out_cpu);
	free(h_out_gpu);
	free(h_out_gpu_shared);
	return 0;
}

